from loguru import logger
from http import HTTPStatus
from fastapi import APIRouter
from code.app.config.exception import ResponseExceptionEnum
from code.app.enum.enums import Model
from code.app.services.api_call import APICall
from code.app.services.response import Response

# 图像处理公共接口路由
image_app = APIRouter()


@image_app.get('/image-generation/result')
def get_result(model: str, task_id: str) -> Response:
    # 检查模型
    if model not in [Model.wanx_style_repaint_v1.model, Model.image_out_painting.model]:
        return Response.fail(error=ResponseExceptionEnum.MODEL_NOT_SUPPORT, data=None)

    # 获取结果
    api_response = APICall.get_result(task_id=task_id)
    logger.info(f'image-generation get_result api_response：{api_response.__str__()}')

    # 按模型获取结果
    if api_response.status_code != HTTPStatus.OK or (api_response.message and api_response.message.strip()):
        error = ResponseExceptionEnum.API_CALL_ERROR
        if api_response.message is not None:
            error.modify_message(api_response.message)
        response = Response.fail(error=error, data=api_response)
        return response

    if model == Model.wanx_style_repaint_v1.model:
        return Response.success(data=api_response)
    elif model == Model.image_out_painting.model:
        results = []
        output = getattr(api_response, 'output', None)
        if isinstance(output, dict):
            output_image_url = output.get('output_image_url', None)
            if output_image_url:
                results.append({'url': output_image_url})
        api_response.output.results = results
    else:
        logger.info(f'model：{model}，暂不支持')

    return Response.success(data=api_response)
